import os
import json5

from macroeconomic_analysis_agent.datascientist import DataScientistAgent
from macroeconomic_analysis_agent.reporter import ReporterAgent

from macroeconomic_analysis_agent.utils import parse_code, generate_timestamped_random_string
from base_agent import BaseAgent

planner_prompt = """
你是一名规划者，你的目标是理解以下每一个智能代理及其能力，使用提供的资源，通过协调智能代理来解决用户的问题。
请利用你的LLM的知识和理解能力来理解用户问题的意图和目标，生成一个可以通过智能代理协作完成而无需用户协助的任务计划。
请逐步思考以实现目标。你可以使用下面给出的资源。同时，请严格遵守重要提醒中的约束和规范。

任务目标:
{user_input}

任务分配说明:
DataScientist单次任务可一次执行查询与分析，不需要拆分为单独的查询和分析任务

可用智能体:
{agents}


约束:
描述说明使用中文。
任务计划的每个步骤都应该是为了推进解决用户目标而存在，不要生成无意义的任务步骤，确保每个步骤内目标明确内容完整。
关注任务计划每个步骤的依赖关系和逻辑，被依赖步骤要考虑被依赖的数据，是否能基于当前目标得到，如果不能请在目标中提示要生成被依赖数据。
每个步骤都是一个独立可完成的目标，一定要确保逻辑和信息完整，不要出现类似:'Analyze the retrieved issues data'这样目标不明确，不知道具体要分析啥内容的步骤。
请确保只使用上面提到的智能体，并且可以只使用其中需要的部分，严格根据描述能力和限制分配给合适的步骤，每个智能体都可以重复使用。
根据用户目标的实际需要使用提供的资源来协助生成计划步骤，不要使用不需要的资源。
每个步骤最好只使用一种资源完成一个子目标，如果当前目标可以分解为同类型的多个子任务，可以生成相互不依赖的并行任务。
数据资源可以被合适的智能体加载使用，不用考虑数据资源的加载链接问题。
尽量合并有顺序依赖的连续相同步骤,如果用户目标无拆分必要，可以生成内容为用户目标的单步任务。
仔细检查计划，确保计划完整的包含了用户问题所涉及的所有信息，并且最终能完成目标，确认每个步骤是否包含了需要用到的资源信息,如URL、资源名等。

按照以下格式输出JSON列表:
```json
[
    {{
        "task_id": str = "unique identifier for a task in plan, can be an ordinal",
        "dependent_task_ids": list[str] = "ids of tasks prerequisite to this task",
        "instruction": "what you should do in this task, one short phrase or sentence",
        "task_type": "type of this task, should be one of Available Task Types",
    }},
    ...
]
```

例子:
任务目标：使用2015年至今的数据，从年度gdp增长率、年度gdp规模生成一份宏观经济报告
输出:
```json
[
    {{
        "task_id": "1",
        "dependent_task_ids": [],
        "instruction": "从数据库中查询2015年至今的年度GDP增长率数据,并对数据进行分析",
        "task_type": "DataScientist"
    }},
    {{
        "task_id": "2",
        "dependent_task_ids": [],
        "instruction": "从数据库中查询2015年至今的年度GDP规模数据，并对数据进行分析",
        "task_type": "DataScientist"
    }},
    {{
        "task_id": "3",
        "dependent_task_ids": ["1", "2"],
        "instruction": "构建包含2015年至今年度GDP增长率与GDP规模分析的专业宏观经济报告",
        "task_type": "Reporter"
    }}
]
```
"""

agents = """
- **DataScientist**: 使用数据库资源生成SQL并执行，根据查询出得数据进行分析。
- **Reporter**: 观察并组织各种分析结果，构建专业报告。
"""

user_input = """

- **需要图表展示**：近二十年GDP增长率按季度变化情况。
SELECT CONCAT(`year`, '-', LPAD(`month`, 2, '0')) AS '日期', `data_value` as '国内生产总值同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%国内生产总值同比增长%' AND `groupNameCn` LIKE '%季度GDP%' AND `year` >= '1995' order by CAST(`year` AS UNSIGNED) ASC, CAST(`month` AS UNSIGNED) ASC约120条记录

- **需要图表展示**：按月统计近五年M0增速的变化情况
M0约108条记录：
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%流通中现金(M0)%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '流通中现金(M0)',MAX(CASE WHEN norm LIKE '%流通中现金(M0)同比增长%' THEN data_value END) AS '流通中现金(M0)同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M0%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近五年M1增速的变化情况
M1约108条记录：
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%货币(狭义货币M1)%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '货币(狭义货币M1)',MAX(CASE WHEN norm LIKE '%货币(狭义货币M1)同比增长%' THEN data_value END) AS '货币(狭义货币M1)同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M1%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近五年M2增速的变化情况
M2约108条记录:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%货币和准货币（广义货币M2）%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '货币和准货币（广义货币M2）',MAX(CASE WHEN norm LIKE '%货币和准货币（广义货币M2）同比增长%' THEN data_value END) AS '货币和准货币（广义货币M2）同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M2%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近10年CPI指数变化情况。
CPI约114条记录
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期', `data_value` as 'CPI居民消费价格指数' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%价格指数%' AND `groupNameCn` LIKE '%CPI居民消费价格指数%' AND `year` >= '2015' ORDER BY CAST(`year` AS UNSIGNED) ASC, CAST(`month` AS UNSIGNED) ASC;

- **无特殊图表展示要求**：提供2015年来进出口总额数据及变化情况，进口与出口产品分类数据及变化情况。
进出口约54条记录
SELECT YEAR AS '日期',MAX(CASE WHEN norm LIKE '%一般贸易出口%' THEN data_value ELSE NULL END) AS '一般贸易出口的值',MAX(CASE WHEN norm LIKE '%一般贸易进口%' THEN data_value ELSE NULL END) AS '一般贸易进口',MAX(CASE WHEN norm LIKE '%加工贸易出口%' THEN data_value ELSE NULL END) AS '加工贸易出口',MAX(CASE WHEN norm LIKE '%加工贸易进口%' THEN data_value ELSE NULL END) AS '加工贸易进口',MAX(CASE WHEN norm LIKE '%其他贸易出口%' THEN data_value ELSE NULL END) AS '其他贸易出口',MAX(CASE WHEN norm LIKE '%其他贸易进口%' THEN data_value ELSE NULL END) AS '其他贸易进口' FROM `data`.`sina_data_2` WHERE `groupNameCn` LIKE '%按贸易分进出口总额%' AND `year`>='2015' GROUP BY YEAR ORDER BY CAST(YEAR AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近五年投资完成额变化图。
约49条记录
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm='投资完成额' THEN data_value ELSE NULL END) AS '投资完成额的值',MAX(CASE WHEN norm='投资完成额同比增长' THEN data_value ELSE NULL END) AS '投资完成额同比增长的值' FROM `data`.`sina_data_2` WHERE `norm` IN ('投资完成额','投资完成额同比增长') AND `groupNameCn` LIKE '%固定资产投资%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近五年建筑业总产值增长率变化情况。
约4条
SELECT `YEAR` AS '日期', `data_value` as '建筑业总产值' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%建筑业%' AND `year` >= '2020' AND `groupNameCn` = '年度GDP' ORDER BY CAST(`year` AS UNSIGNED) ASC;

- **需要图表展示**：按月统计近五年PMI指数、新出口订单指数、新订单指数。
约162条记录
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm='PMI' THEN data_value ELSE NULL END) AS 'PMI',MAX(CASE WHEN norm='新订单指数' THEN data_value ELSE NULL END) AS '新订单指数',MAX(CASE WHEN norm='新出口订单指数' THEN data_value ELSE NULL END) AS '新出口订单指数' FROM `data`.`sina_data_2` WHERE `norm` IN ('PMI','新订单指数','新出口订单指数') AND `year`>='2020' AND `groupNameCn` LIKE '%制造业PMI%' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;


根据以上内容中列出的各项指标生成一份宏观经济报告
"""


def init_agents():
    agents = """
    - **DataScientist**: 使用数据库资源生成SQL并执行，根据查询出得数据进行分析。
    - **Reporter**: 观察并组织各种分析结果，构建专业报告。
    """
    return agents


class MacroeconomicAnalysisAgent(BaseAgent):

    def __init__(self, model, user_input, **kwargs):
        """
        类的初始化方法.

        参数:
        - model: 模型实例，用于处理数据和执行任务.
        - user_input: 用户输入，用于定制计划和任务.
        - **kwargs: 其他关键字参数，用于提供额外的配置.

        返回:
        无返回值.
        """
        super().__init__(model=model, **kwargs)  # 调用基类构造方法
        self.plan_id = generate_timestamped_random_string()  # 生成带有时间戳的随机字符串作为计划ID
        self.agents = init_agents()  # 初始化代理对象，用于处理特定任务
        self.agent_list = {"DataScientist": DataScientistAgent, "Reporter": ReporterAgent}  # 定义可用的代理类型
        self.user_input = user_input  # 保存用户输入，以便后续使用
        self.plan = None  # 初始化计划，将在后续步骤中生成
        self.planner_prompt = planner_prompt.format(user_input=self.user_input, agents=self.agents)  # 根据用户输入和代理生成规划提示

    def create_plan(self):
        """
        创建执行计划。

        通过调用语言模型生成任务计划。此方法会尝试多次调用语言模型，
        直到成功或达到最大尝试次数。生成的任务计划将被解析并存储在
        实例的 plan 属性中。
        """
        # 初始化语言模型调用成功标志为False
        call_llm_success = False
        # 初始化语言模型调用次数计数器
        call_llm_count = 0
        # 初始化任务文本为空字符串
        tasks_text = ''
        # 构建发送给语言模型的消息，当前只有一个用户角色的消息
        messages = [{'role': 'user', 'content': self.planner_prompt}]

        # 循环尝试调用语言模型，直到成功或达到最大尝试次数（10次）
        while not call_llm_success and call_llm_count < 10:
            # 调用语言模型并获取响应
            resp = self._call_llm(messages=messages)
            # 重置任务文本
            tasks_text = ''
            # 遍历响应内容并累加到任务文本中
            for r in resp:
                tasks_text += r
            # 检查任务文本中是否包含错误代码
            if 'Error code' in tasks_text:
                # 如果包含错误代码，则增加调用尝试次数
                call_llm_count += 1
            else:
                # 如果不包含错误代码，则设置成功标志为True，表示调用成功
                call_llm_success = True

        # 打印任务文本
        print('Tasks_text: ', tasks_text)
        # 将任务文本解析为指定语言的代码。这里假设为JSON
        tasks_text = parse_code(text=tasks_text, lang='json')
        # 将解析后的任务文本解析为JSON对象
        tasks = json5.loads(tasks_text)
        # 将解析后的任务计划存储在实例的plan属性中
        self.plan = tasks

    def execute_plan(self, task):
        """
        根据任务类型执行相应代理的计划。

        从agent_list中选择与任务类型相匹配的代理类，创建该代理类的实例并调用其run方法，
        执行任务，并返回任务执行结果。

        参数:
        - task (dict): 包含任务类型和任务具体信息的任务字典。

        返回:
        - is_success (bool): 任务执行结果，成功返回True，失败返回False。
        """
        # 根据任务类型选择合适的代理类
        Agent = self.agent_list[task['task_type']]
        # 创建代理实例并调用run方法执行任务
        is_success = Agent(task, self.model, self.plan_id).run()
        # 返回任务执行结果
        return is_success

    def create_directory(self):
        os.makedirs(self.plan_id, exist_ok=True)

    def run(self):
        """
        执行计划的主要方法。
        首先创建计划，然后逐个执行计划中的任务。
        如果某个任务执行失败且在允许的重试次数内未能成功，则记录失败信息并终止执行。
        如果所有任务都成功执行，则返回成功标志和完成信息。

        返回:
            元组，包含一个布尔值表示执行是否成功和相关信息字符串。
        """
        # 创建执行计划
        self.create_plan()
        self.create_directory()
        # 遍历计划中的每个任务
        for task in self.plan:
            print(f"{task['task_type']}：task_id为{task['task_id']}的任务{task['instruction']}开始执行")
            is_success = False
            max_retry = 0

            # 尝试执行任务，直到成功或达到最大重试次数
            while max_retry < 5:
                result = self.execute_plan(task)
                if result:
                    is_success = True
                    break
                else:
                    max_retry += 1

            # 如果任务最终未能成功执行，返回失败信息
            if not is_success:
                return False, f"{task['task_type']}：task_id为{task['task_id']}的任务{task['instruction']}执行失败"
            print(f"{task['task_type']}：task_id为{task['task_id']}的任务{task['instruction']}执行完成")
        print(f'所有任务执行完成, 任务id：{self.plan_id}')
        # 所有任务成功执行后，返回成功标志和完成信息
        return True, '所有任务执行完成，请到指定目录查看报告'


if __name__ == '__main__':
    agent = MacroeconomicAnalysisAgent(model='qwen2-72b-instruct', user_input=user_input)
    a, b = agent.run()
    print(a)
    print(b)
